"""9.3.1 节代码：tutor_agent.py

Small script that composes a LangChain agent (Qianfan LLM + Bocha tool)
and runs an interactive REPL. The file was reformatted for readability while
preserving original behavior.
"""

import os
import time

from langchain_community.llms import QianfanLLMEndpoint
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.prompts import PromptTemplate  # 修改这里：使用 PromptTemplate 而不是 ChatPromptTemplate
from langchain_core.messages import SystemMessage

from bocha_tool import get_bocha_search_tool


# --- Part 1: Setup & Components ---
# 1. 安全加载 API 密钥 (百度千帆)
os.environ["QIANFAN_ACCESS_KEY"] = os.getenv("QIANFAN_ACCESS_KEY")
os.environ["QIANFAN_SECRET_KEY"] = os.getenv("QIANFAN_SECRET_KEY")

# 2. 积木：Model (AI 的大脑)
llm = QianfanLLMEndpoint(model="ERNIE-Bot-4")

# 3. 积木：Memory (AI 的记事本)
# return_messages=True 是 Agent 框架的要求
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 4. 积木：Tools (AI 的工具箱)
tools = [get_bocha_search_tool()]

# 5. 积木：System Prompt (AI 的人设和决策指令)
# 这是 Agent 的核心，它定义了 Alex 的角色和使用工具的规则
system_prompt_content = (
    "You are a friendly, patient, and encouraging English tutor named Alex.\n"
    "Your goal is to help a non-native speaker practice conversational English.\n"
    "Always ask a follow-up question to keep the conversation going.\n"
    "Keep your responses concise, generally one or two sentences.\n"
    "If the user asks a question that requires up-to-date, real-time information "
    "(like current news, sports scores, or today's weather), you MUST use the 'Bing "
    "Search' tool.\n"
    "If the user makes a grammatical error, gently correct it and explain the "
    "rule, but DO NOT use the search tool for corrections.\n\n"
    "{tools}\n\n"  # 添加工具描述
    "Use the following format:\n\n"
    "Question: the input question you must answer\n"
    "Thought: you should always think about what to do\n"
    "Action: the action to take, should be one of [{tool_names}]\n"
    "Action Input: the input to the action\n"
    "Observation: the result of the action\n"
    "... (this Thought/Action/Action Input/Observation can repeat N times)\n"
    "Thought: I now know the final answer\n"
    "Final Answer: the final answer to the original input question\n\n"
    "Begin!\n\n"
    "Previous conversation:\n{chat_history}\n\n"  # 添加对话历史
    "Question: {input}\n"
    "Thought:{agent_scratchpad}"
)


# --- Part 2: 构建 Agent ---
# 1. 创建 Agent Prompt Template
# 使用 PromptTemplate 而不是 ChatPromptTemplate
prompt = PromptTemplate(
    template=system_prompt_content,
    input_variables=["tools", "tool_names", "input", "chat_history", "agent_scratchpad"]
)

# 2. 创建 Agent 实例 (使用 ReAct 框架)
# ReAct 是一种流行的 Agent 框架，它指导 LLM 进行 "Reasoning" (思考) 和 "Action" (行动)
agent = create_react_agent(llm, tools, prompt)

# 3. 创建 Agent Executor (执行器)
# Executor 负责运行 Agent 的推理循环，直到得出最终答案
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,  # 开启 verbose，让我们看到 Agent 的思考过程
    handle_parsing_errors=True,
)


# --- Part 3: 对话循环 ---
if __name__ == "__main__":
    print(
        "Hello! I am Alex, your full-stack AI English tutor. I can even search the web for you!"
    )
    print(
        "Try asking me about a current event, or make a grammar mistake. Type 'exit' when done."
    )

    while True:
        user_input = input("\nYou: ")
        if user_input.lower() == "exit":
            print("Alex: Goodbye! It was great practicing with you.")
            break

        try:
            print("Alex is thinking...")
            start_time = time.time()

            # 关键：调用 Agent Executor，而不是 Chain
            response = agent_executor.invoke({"input": user_input})

            end_time = time.time()
            print(f"\n思考总耗时: {end_time - start_time:.2f} 秒")

            # Agent Executor 的最终回答在 output 键中
            print("Alex: " + response["output"])
        except Exception as e:
            print(
                f"Alex: I encountered an error while processing your request. Error: {e}"
            )